import threading
import numpy as np
np.random.seed(10)
import torch
class Buffer:
    def __init__(self, args):
        self.size = args.buffer_size
        self.args = args
        self.current_size = 0
        self.sample_count = 0
        
        # 创建缓冲区，观测值维度改为10
        self.buffer = dict()
        for i in range(self.args.num_uavs):
            self.buffer['o_%d' % i] = np.empty([self.size, 10])  # 简化为10维
            self.buffer['u_%d' % i] = np.empty([self.size, 4])
            self.buffer['r_%d' % i] = np.empty([self.size])
            self.buffer['o_next_%d' % i] = np.empty([self.size, 10])  # 简化为10维
        
        # IB缓冲区保持不变
        self.IB_Buffer = dict()
        self.IB_Buffer['info'] = np.empty([self.size, 8])
        self.IB_Buffer['mu'] = np.empty([self.size, 8])
        self.IB_Buffer['logvar'] = np.empty([self.size, 8])
        
        # 线程锁
        self.lock = threading.Lock()
    
    def store_episode(self, o, u, r, o_next, agent_id):
        """存储单个智能体的经验"""
        import torch
        import numpy as np
        
        idxs = self._get_storage_idx(inc=1)
        
        with self.lock:
            # 处理观测值
            if isinstance(o, torch.Tensor):
                self.buffer['o_%d' % agent_id][idxs] = o.detach().numpy()
            elif isinstance(o, np.ndarray):
                self.buffer['o_%d' % agent_id][idxs] = o
            else:
                self.buffer['o_%d' % agent_id][idxs] = np.array(o)
            
            # 处理动作
            if isinstance(u, torch.Tensor):
                self.buffer['u_%d' % agent_id][idxs] = u.detach().numpy()
            elif isinstance(u, np.ndarray):
                self.buffer['u_%d' % agent_id][idxs] = u
            else:
                self.buffer['u_%d' % agent_id][idxs] = np.array(u)
            
            # 处理奖励
            if isinstance(r, torch.Tensor):
                self.buffer['r_%d' % agent_id][idxs] = r.detach().numpy()
            elif isinstance(r, np.ndarray):
                self.buffer['r_%d' % agent_id][idxs] = r
            else:
                self.buffer['r_%d' % agent_id][idxs] = np.array(r)
            
            # 处理下一个观测值
            if isinstance(o_next, torch.Tensor):
                self.buffer['o_next_%d' % agent_id][idxs] = o_next.detach().numpy()
            elif isinstance(o_next, np.ndarray):
                self.buffer['o_next_%d' % agent_id][idxs] = o_next
            else:
                self.buffer['o_next_%d' % agent_id][idxs] = np.array(o_next)
    
                
                
    
    
    # sample the data from the replay buffer
    def sample(self, batch_size):
        temp_buffer = {}
        if self.sample_count %3 == 0:
            np.random.seed(self.sample_count)
        idx = np.random.randint(0, self.current_size, batch_size)
        for key in self.buffer.keys():
            temp_buffer[key] = self.buffer[key][idx]
        return temp_buffer

    def _get_storage_idx(self, inc=None):
        inc = inc or 1
        if self.current_size+inc <= self.size:
            idx = np.arange(self.current_size, self.current_size+inc)
        elif self.current_size < self.size:
            overflow = inc - (self.size - self.current_size)
            idx_a = np.arange(self.current_size, self.size)
            idx_b = np.random.randint(0, self.current_size, overflow)
            idx = np.concatenate([idx_a, idx_b])
        else:
            idx = np.random.randint(0, self.size, inc)
        self.current_size = min(self.size, self.current_size+inc)
        if inc == 1:
            idx = idx[0]
        return idx
